# coding:utf-8
import joblib
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from tqdm import tqdm

# 导入绘图相关库
import matplotlib.pyplot as plt
from sklearn.metrics import roc_curve, auc

def createModel():
    # 初始化进度条（总进度100%，中文提示）
    with tqdm(total=100, desc="模型生成进度") as pbar:
        # 1. 读取数据（占比10%）
        pbar.set_description("正在读取数据...")
        names = ['Sample code number', 'Clump Thickness', 'Uniformity of Cell Size', 'Uniformity of Cell Shape',
                 'Marginal Adhesion', 'Single Epithelial Cell Size', 'Bare Nuclei', 'Bland Chromatin',
                 'Normal Nucleoli', 'Mitoses', 'Class']
        data = pd.read_csv(
            'https://archive.ics.uci.edu/ml/machine-learning-databases/breast-cancer-wisconsin/breast-cancer-wisconsin.data',
            names=names)
        print(data.head())
        pbar.update(10)  # 进度+10%

        # 2. 数据清洗（占比10%）
        pbar.set_description("正在清洗数据...")
        data = data.replace(to_replace='?', value=np.nan).dropna()
        x = data.iloc[:, 1:10]
        y = data['Class']
        pbar.update(10)  # 进度+10%（累计20%）

        # 3. 划分训练测试集（占比10%）
        pbar.set_description("划分训练测试集...")
        x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=22)
        pbar.update(10)  # 进度+10%（累计30%）

        # 4. 数据标准化（占比20%）
        pbar.set_description("标准化数据...")
        transfer = StandardScaler()
        x_train = transfer.fit_transform(x_train)
        x_test = transfer.transform(x_test)
        pbar.update(20)  # 进度+20%（累计50%）

        # 5. 模型训练（占比30%）
        pbar.set_description("模型训练中...")
        estimator = LogisticRegression()
        estimator.fit(x_train, y_train)
        pbar.update(30)  # 进度+30%（累计80%）

        # 6. 模型评估（占比10%）
        pbar.set_description("评估模型效果...")
        y_pred = estimator.predict(x_test)
        print("预测值是:\n", y_pred)
        print("测试集的准确率是:", estimator.score(x_test, y_test))

        # 预测概率
        y_scores = estimator.predict_proba(x_test)[:, 1]  # 获取正类（恶性）的概率

        # 计算 ROC 曲线指标
        fpr, tpr, thresholds = roc_curve(y_test, y_scores, pos_label=4)  # 标签4表示恶性
        roc_auc = auc(fpr, tpr)

        # 绘制 ROC 曲线
        plt.figure(figsize=(8, 6))
        plt.plot(fpr, tpr, color='blue', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
        plt.plot([0, 1], [0, 1], color='gray', linestyle='--', lw=2)
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('False Positive Rate')
        plt.ylabel('True Positive Rate')
        plt.title('Receiver Operating Characteristic (ROC) Curve')
        plt.legend(loc="lower right")
        plt.show()

        pbar.update(10)  # 进度+10%（累计90%）

        # 7. 保存模型（占比10%）
        pbar.set_description("保存模型文件...")
        joblib.dump(estimator, 'logistic_regression_model.pkl')
        joblib.dump(transfer, 'scaler.pkl')
        print("模型和标准化器已保存至当前目录")
        pbar.update(10)  # 进度+10%（累计100%）

    input("请按任意键继续...")


def loadModel():
    try:
        # 加载模型和标准化器（关键修改：从文件加载scaler）
        loaded_model = joblib.load('logistic_regression_model.pkl')
        transfer = joblib.load('scaler.pkl')  # 加载训练时的标准化器
    except FileNotFoundError:
        print("错误：未找到模型文件，请先运行'生成模型'选项")
        return

    # 生成示例测试数据（特征顺序与训练时一致）
    test_samples = pd.DataFrame({
        'Clump Thickness': [2, 5, 8],
        'Uniformity of Cell Size': [1, 3, 6],
        'Uniformity of Cell Shape': [1, 2, 5],
        'Marginal Adhesion': [1, 2, 4],
        'Single Epithelial Cell Size': [2, 3, 4],
        'Bare Nuclei': [1, 2, 7],
        'Bland Chromatin': [2, 4, 6],
        'Normal Nucleoli': [1, 2, 3],
        'Mitoses': [1, 1, 2]
    })

    # 标准化测试数据（使用训练时的scaler）
    test_samples_scaled = transfer.transform(test_samples)

    # 模型预测并生成表格
    test_pred = loaded_model.predict(test_samples_scaled)
    result_table = test_samples.copy()
    result_table['预测类别'] = test_pred
    print("\n===== 测试数据预测结果表格 =====")
    print(result_table)
    input("请按任意键继续...")

def user_input_predict():
    """用户手动输入特征数据并预测"""
    try:
        # 加载模型和标准化器（与loadModel共用）
        loaded_model = joblib.load('logistic_regression_model.pkl')
        transfer = joblib.load('scaler.pkl')
    except FileNotFoundError:
        print("错误：未找到模型文件，请先运行'生成模型'选项")
        return

    # 定义特征名称（顺序必须与训练时一致）
    features = [
        'Clump Thickness',
        'Uniformity of Cell Size',
        'Uniformity of Cell Shape',
        'Marginal Adhesion',
        'Single Epithelial Cell Size',
        'Bare Nuclei',
        'Bland Chromatin',
        'Normal Nucleoli',
        'Mitoses'
    ]

    # 用户输入并验证（循环直到输入有效）
    input_data = {}
    for feature in features:
        while True:
            try:
                value = int(input(f"请输入 {feature}（1-10的整数）: "))
                if 1 <= value <= 10:  # 原数据特征值范围1-10
                    input_data[feature] = [value]  # 转为列表适配DataFrame
                    break
                else:
                    print("输入值需在1到10之间，请重新输入")
            except ValueError:
                print("输入错误，请输入整数")

    # 数据标准化并预测
    user_df = pd.DataFrame(input_data)
    user_scaled = transfer.transform(user_df)
    pred = loaded_model.predict(user_scaled)

    # 转换为中文标签
    class_mapping = {2: '良性', 4: '恶性'}
    print(f"\n预测结果：该乳腺肿瘤样本为 {class_mapping[pred[0]]}")
    input("请按任意键继续...")

# 菜单函数修改（新增选项3）
def main_menu():
    while True:
        print("\n===== 乳腺癌分类模型系统 =====")
        print("1. 生成模型（训练并保存）")
        print("2. 运行模型（加载并预测示例数据）")
        print("3. 用户输入预测（手动输入数据预测）")  # 新增选项
        print("4. 退出系统")
        choice = input("请输入选项（1/2/3/4）: ").strip()

        if choice == '1':
            createModel()
        elif choice == '2':
            loadModel()
        elif choice == '3':
            user_input_predict()  # 调用新函数
        elif choice == '4':
            print("系统已退出，感谢使用！")
            break
        else:
            print("输入错误，请输入1-4的数字")

# ... 启动菜单代码保持不变 ...

# 启动菜单
if __name__ == "__main__":
    main_menu()